package com.djf.springaidemo.config;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.preretrieval.query.transformation.CompressionQueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.TranslationQueryTransformer;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

/**
 * 配置类用于设置RAG（Retrieval-Augmented Generation）模型的相关组件
 */
@Configuration
public  class RagConfig {

    /**
     * 创建一个VectorStore实例，用于存储和查询文档的向量表示
     *
     * @param embeddingModel 用于生成文档向量的嵌入模型
     * @return VectorStore的实例
     */
    @Bean
    public VectorStore vectorStore2(EmbeddingModel embeddingModel) {
        SimpleVectorStore vectorStore = SimpleVectorStore.builder(embeddingModel)
                .build();

        // 初始化示例文档
        vectorStore.add(List.of(
                new Document("产品说明书:产品名称：智能机器人\n" +
                        "产品描述：智能机器人是一个智能设备，能够自动完成各种任务。\n" +
                        "功能：\n" +
                        "1. 自动导航：机器人能够自动导航到指定位置。\n" +
                        "2. 自动抓取：机器人能够自动抓取物品。\n" +
                        "3. 自动放置：机器人能够自动放置物品。\n")
        ));
        return vectorStore;
    }

    /**
     * 创建一个MultiQueryExpander实例，用于扩展查询以提高检索效果
     *
     * @param builder 用于构建ChatClient的构建器
     * @return MultiQueryExpander的实例
     */
    @Bean
    public MultiQueryExpander multiQueryExpander(ChatClient.Builder builder) {
        return MultiQueryExpander.builder()
                .chatClientBuilder(builder)
                .includeOriginal(false)
                .numberOfQueries(3)
                .build();
    }

    /**
     * 创建一个RewriteQueryTransformer实例，用于重写查询以提高检索效果
     *
     * @param builder 用于构建ChatClient的构建器
     * @return RewriteQueryTransformer的实例
     */
    @Bean
    public RewriteQueryTransformer rewriteQueryTransformer(ChatClient.Builder builder) {
        return RewriteQueryTransformer.builder()
                .chatClientBuilder(builder)
                .build();
    }

    /**
     * 创建一个TranslationQueryTransformer实例，用于将查询翻译成目标语言以提高检索效果
     *
     * @param builder 用于构建ChatClient的构建器
     * @return TranslationQueryTransformer的实例
     */
    @Bean
    public TranslationQueryTransformer translationQueryTransformer(ChatClient.Builder builder) {
        return TranslationQueryTransformer.builder()
                .chatClientBuilder(builder)
                .targetLanguage("chinese")
                .build();
    }

    /**
     * 创建一个CompressionQueryTransformer实例，用于压缩查询以提高检索效果
     *
     * @param builder 用于构建ChatClient的构建器
     * @return CompressionQueryTransformer的实例
     */
    @Bean
    public CompressionQueryTransformer compressionQueryTransformer(ChatClient.Builder builder) {
        return CompressionQueryTransformer.builder()
                .chatClientBuilder(builder)
                .build();
    }

    /**
     * 创建一个RetrievalAugmentationAdvisor实例，用于指导RAG模型的检索增强过程
     *
     * @param retriever 用于检索文档的DocumentRetriever实例
     * @return RetrievalAugmentationAdvisor的实例
     */
    @Bean
    public RetrievalAugmentationAdvisor retrievalAdvisor(DocumentRetriever retriever) {
        return RetrievalAugmentationAdvisor.builder()
                .documentRetriever(retriever)
                .queryAugmenter(ContextualQueryAugmenter.builder()
                        .allowEmptyContext(true)
                        .build())
                .build();
    }

    /**
     * 创建一个DocumentRetriever实例，用于从VectorStore中检索相关文档
     *
     * @param vectorStore 用于存储文档向量的VectorStore实例
     * @return DocumentRetriever的实例
     */
    @Bean
    public DocumentRetriever documentRetriever(VectorStore vectorStore) {
        return VectorStoreDocumentRetriever.builder()
                .vectorStore(vectorStore)
                .similarityThreshold(0.5)
                .topK(3)
                .build();
    }

}
